#!/usr/bin/env python
# coding=utf-8
'''
Author: 以谁为师
Website: attacker.club
Date: 2020-11-16 21:00:25
LastEditTime: 2020-11-17 00:47:01
Description:
'''
import argparse
import subprocess
import paramiko
import socket
import time
import sys
import os
import re
# import logging
# from logging.handlers import RotatingFileHandler
from log import LoggerLog
logger = LoggerLog("ssh").get_log

results = {
    'unsuccessful': [],
    'successful': [],
}


def parse_args():

    content = """ssh远程工具 ^_^

    远程执行:  ssh_paramiko.py  -h 192.168.0.111  -uroot -p123456  --run "who -b"
              ssh_paramiko.py  -h 192.168.0.111  -uroot -k ~/.ssh/id_rsa  -P55555 --run "who -b"

    上传:   ssh_paramiko.py  -h attacker.club  -uroot -k ~/.ssh/id_rsa  -P55555  -t put  --local_file log.py   --remote_dir /tmp/
    下载:   ./ssh_paramiko.py  -h attacker.club  -uroot -k ~/.ssh/id_rsa  -P55555  -t get  --file iptables.init    --remote_dir /opt  --local_dir /tmp
        """
    parser = argparse.ArgumentParser(
        usage="paramiko ssh",
        description=content,
        add_help=False,
        formatter_class=lambda prog: argparse.RawTextHelpFormatter(
            prog, max_help_position=50)
    )
    parser.add_argument("--help",
                        action="help",
                        help="查看帮助信息")
    parser.add_argument("-t", "--type",
                        default="ssh",
                        help="ssh,put,get ;默认使用ssh")

    parser.add_argument('-u', '--user',
                        default="root",
                        help='系统用户名,默认root')

    parser.add_argument('-h', '--host',
                        help='服务器IP')

    parser.add_argument('-p', '--passwd',
                        help='服务器密码')
    parser.add_argument('-k', '--key',
                        help='SSH key路径')

    parser.add_argument('-P', '--port',
                        type=int,
                        default=22,
                        help='服务器端口,默认为22')

    parser.add_argument('-r', '--run',
                        default="uptime",
                        help='执行命令,默认 uptime')
    parser.add_argument('--file',
                        help='文件名')
    parser.add_argument('--local_file',
                        help='本地文件')
    parser.add_argument('--local_dir',
                        help='本地目录')
    parser.add_argument('--remote_dir',
                        help='远程目录')

    args = parser.parse_args()
    return args


class SSHConnection(object):

    log_path = os.path.dirname(os.path.abspath(__file__))
    os.makedirs('logs', exist_ok=True)  # 设置日志输出目录
    logname = log_path + '/logs/paramiko.log'  # 指定输出的日志文件名
    paramiko.util.log_to_file(logname)

    def __init__(self, args):
        self.con = args

    def ssh_connect(self):
        logger.info(
            "[\033[1;32m%s\033[0m] Trying to connect　..." % (self.con.host))
        self.sshclient = paramiko.SSHClient()
        self.sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            if self.con.passwd:
                self.sshclient.connect(
                    hostname=self.con.host,
                    username=self.con.user,
                    password=self.con.passwd,
                    port=self.con.port,
                    timeout=3, allow_agent=False)
            else:

                private_key = paramiko.RSAKey.from_private_key_file(
                    self.con.key)
                self.sshclient.connect(
                    hostname=self.con.host,
                    username=self.con.user,
                    pkey=private_key,
                    port=self.con.port,
                    timeout=3, allow_agent=False, look_for_keys=False)

        except (paramiko.ssh_exception.AuthenticationException,) as e:
            logger.error("\033[1;31mLogin failed\033[0m\t%s" % (e))
            exit(1)
        except (paramiko.ssh_exception.NoValidConnectionsError, socket.gaierror, socket.timeout, socket.error) as e:
            logger.error("\033[1;31mConnection failed\033[0m\t%s" % (e))
            exit(1)
        except Exception as e:
            logger.error(e, type(e))
            exit(1)

    def remote_exec(self):
        date = time.strftime("%b %d %H:%M:%S", time.localtime())
        stdin, stdout, stderr = self.sshclient.exec_command(
            "source /etc/profile ; %s" % self.con.run)
        result = stdout.read().decode() + stderr.read().decode()
        return result

    def remote_put_file(self):
        try:
            logger.info(
                "[\033[1;32m%s\033[0m] to uploading　..." % (self.con.local_file))
            transport = self.sshclient.get_transport()
            sftp = paramiko.SFTPClient.from_transport(transport)
            remote_file = os.path.join(
                self.con.remote_dir, os.path.basename(self.con.local_file))
            sftp.put(self.con.local_file, remote_file)

            logger.info('上传文件成功：%s' % self.con.local_file)
        except Exception:
            logger.error('上传文件失败：%s' % self.con.local_file)
            sys.exit(1)

    def sftp_get_file(self):
        try:
            logger.info(
                "[\033[1;32m%s\033[0m] to downloading　..." % (self.con.file))
            transport = self.sshclient.get_transport()
            sftp = paramiko.SFTPClient.from_transport(transport)
            rfile = os.path.join(self.con.remote_dir, self.con.file)
            lfile = os.path.join(self.con.local_dir, self.con.file)
            # print(rfile, lfile)
            sftp.get(rfile, lfile)
            logger.info('下载文件成功：%s' % self.con.file)
        except Exception as e:
            logger.error('下载文件失败：%s %s' % (self.con.file, e))
            sys.exit(1)

    def close_ssh(self):
        try:
            self.sshclient.close()
        except Exception as e:
            logger.error(e)

    @property
    def execute_command(self):
        if args.type == "ssh":
            ret = self.remote_exec()
            print(ret)  # 打印执行结果
        elif args.type == "put":
            ret = self.remote_put_file()
        elif args.type == "get":
            ret = self.sftp_get_file()

        else:
            logger.error("参数错误")

        con.close_ssh()
        logger.info("ssh远程执行结束")


if __name__ == '__main__':
    if len(sys.argv) == 1:
        """ 如果没有传参则打印帮助信息 """
        sys.argv.append("--help")
        args = parse_args()
        sys.exit(1)

    args = parse_args()  # 获取参数

    con = SSHConnection(args)
    con.ssh_connect()
    con.execute_command
